use std::mem::size_of;

use bytemuck::cast_slice;
use wgpu::{include_wgsl, util::DeviceExt};

use crate::Color;

pub struct Renderer2D {
    pub camera: HeightCamera,
    pub camera_buffer: wgpu::Buffer,
    camera_bind_group: wgpu::BindGroup,

    pub rects: Vec<Rect>,
    rect_buffer: wgpu::Buffer,
    rect_fill_pipeline: wgpu::RenderPipeline,

    pub ellipses: Vec<Rect>,
    ellipse_buffer: wgpu::Buffer,
    ellipse_fill_pipeline: wgpu::RenderPipeline,
}

impl Renderer2D {
    /// aspect is the surface_width / surface_height
    pub fn new(
        device: &wgpu::Device,
        aspect: f32,
        surface_format: wgpu::TextureFormat
    ) -> Self {
        let camera = HeightCamera::new([0.0, 0.0], 13.0);
        let camera_proj = camera.to_gpu_proj(aspect);
        let camera_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("camera"),
            contents: cast_slice(&[camera_proj]),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST
        });
        let camera_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("camera"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    visibility: wgpu::ShaderStages::VERTEX,
                    count: None,
                }
            ]
        });
        let camera_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("camera"),
            layout: &camera_bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: camera_buffer.as_entire_binding()
                }
            ]
        });

        let rect_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("rect"),
            size: 0,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let rect_fill_shader = include_wgsl!("shaders/rect_fill.wgsl");
        let rect_fill_module = device.create_shader_module(rect_fill_shader);
        let rect_fill_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("rect fill"),
            bind_group_layouts: &[&camera_bind_group_layout],
            push_constant_ranges: &[]
        });
        let rect_fill_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("rect fill"),
            layout: Some(&rect_fill_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &rect_fill_module,
                entry_point: "vs_main",
                compilation_options: Default::default(),
                buffers: &[
                    wgpu::VertexBufferLayout {
                        array_stride: size_of::<Rect>() as u64,
                        step_mode: wgpu::VertexStepMode::Instance,
                        attributes: &[
                            wgpu::VertexAttribute {
                                format: wgpu::VertexFormat::Float32x2,
                                offset: 0,
                                shader_location: 0
                            },
                            wgpu::VertexAttribute {
                                format: wgpu::VertexFormat::Float32x2,
                                offset: 8,
                                shader_location: 1,
                            },
                            wgpu::VertexAttribute {
                                format: wgpu::VertexFormat::Uint32,
                                offset: 16,
                                shader_location: 2,
                            }
                        ]
                    }
                ]
            },
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleStrip,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: Default::default(),
            fragment: Some(wgpu::FragmentState {
                module: &rect_fill_module,
                entry_point: "fs_main",
                compilation_options: Default::default(),
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })]
            }),
            multiview: None,
            cache: Default::default(),
        });

        let ellipse_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("ellipse"),
            size: 0,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let ellipse_fill_shader = include_wgsl!("shaders/ellipse_fill.wgsl");
        let ellipse_fill_module = device.create_shader_module(ellipse_fill_shader);
        let ellipse_fill_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("ellipse fill"),
            bind_group_layouts: &[&camera_bind_group_layout],
            push_constant_ranges: &[]
        });
        let ellipse_fill_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("ellipse fill"),
            layout: Some(&ellipse_fill_pipeline_layout),
            cache: Default::default(),
            vertex: wgpu::VertexState {
                module: &ellipse_fill_module,
                entry_point: "vs_main",
                compilation_options: Default::default(),
                buffers: &[
                    wgpu::VertexBufferLayout {
                        array_stride: size_of::<Rect>() as u64,
                        step_mode: wgpu::VertexStepMode::Instance,
                        attributes: &[
                            wgpu::VertexAttribute {
                                format: wgpu::VertexFormat::Float32x2,
                                offset: 0,
                                shader_location: 0
                            },
                            wgpu::VertexAttribute {
                                format: wgpu::VertexFormat::Float32x2,
                                offset: 8,
                                shader_location: 1,
                            },
                            wgpu::VertexAttribute {
                                format: wgpu::VertexFormat::Uint32,
                                offset: 16,
                                shader_location: 2,
                            }
                        ]
                    }
                ]
            },
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleStrip,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: Default::default(),
            fragment: Some(wgpu::FragmentState {
                module: &ellipse_fill_module,
                entry_point: "fs_main",
                compilation_options: Default::default(),
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })]
            }),
            multiview: None,
        });
        Self {
            camera, camera_buffer, camera_bind_group,
            rects: Vec::new(), rect_buffer, rect_fill_pipeline,
            ellipses: Vec::new(), ellipse_buffer, ellipse_fill_pipeline,
        }
    }
    pub fn draw<'a>(
        &'a self,
        pass: &mut wgpu::RenderPass<'a>
    ) {
        pass.set_bind_group(0, &self.camera_bind_group, &[]);

        if self.rects.len() > 0 {
            pass.set_vertex_buffer(0, self.rect_buffer.slice(..));
            pass.set_pipeline(&self.rect_fill_pipeline);
            pass.draw(0..4, 0..self.rects.len() as u32);
        }
        if self.ellipses.len() > 0 {
            pass.set_vertex_buffer(0, self.ellipse_buffer.slice(..));
            pass.set_pipeline(&self.ellipse_fill_pipeline);
            pass.draw(0..4, 0..self.ellipses.len() as u32);
        }
    }
    /// update changes to GPU
    pub fn update(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue
    ) {
        let raw = cast_slice(&self.rects);
        if raw.len() as u64 > self.rect_buffer.size() {
            self.rect_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("rect"),
                contents: raw,
                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::VERTEX
            })
        } else {
            queue.write_buffer(&self.rect_buffer, 0, cast_slice(&self.rects));
        }
        let raw = cast_slice(&self.ellipses);
        if raw.len() as u64 > self.ellipse_buffer.size() {
            self.ellipse_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("ellipse"),
                contents: raw,
                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::VERTEX
            })
        } else {
            queue.write_buffer(&self.ellipse_buffer, 0, cast_slice(&self.ellipses));
        }
    }
    /// update aspect when the surface resize
    pub fn resize(&mut self, aspect: f32, queue: &wgpu::Queue) {
        let camera_proj = self.camera.to_gpu_proj(aspect);
        queue.write_buffer(&self.camera_buffer, 0, cast_slice(&[camera_proj]));
    }
}

#[repr(C)]
#[derive(bytemuck::NoUninit, Clone, Copy)]
pub struct CameraProj2d {
    pub pos: [f32; 2],
    pub scale: [f32; 2],
}
#[derive(Clone, Copy)]
pub struct HeightCamera {
    /// world position
    pub pos: [f32; 2],
    /// viewport height in world space
    pub height: f32,
}
impl HeightCamera {
    pub fn new(pos: [f32; 2], height: f32) -> Self {
        Self {
            pos, height,
        }
    }
    /// aspect is the surface_width / surface_height
    /// Transfer the width and height to their invertion ahead for better performance.
    pub fn to_gpu_proj(&self, aspect: f32) -> CameraProj2d {
        let width = self.height * aspect;
        CameraProj2d {
            pos: self.pos,
            scale: [2.0 / width, 2.0 / self.height]
        }
    }
}


#[repr(C)]
#[derive(bytemuck::NoUninit, Clone, Copy)]
pub struct Rect {
    pub min: [f32; 2],
    pub max: [f32; 2],
    pub color: Color,
}
impl Rect {
    pub fn new(min: [f32; 2], max: [f32; 2], color: Color) -> Self {
        Self { min, max, color }
    }
    pub fn from_pos_radius(pos: [f32; 2], radius: f32, color: Color) -> Self {
        Self {
            min: [pos[0] - radius, pos[1] - radius],
            max: [pos[0] + radius, pos[1] + radius],
            color
        }
    }
}